import blue
import yaml
import trinity
METASHADER_PATH = 'res:/Graphics/Shaders/MetashaderLibrary.yaml'
_gMetaShaderLibrary = None
ENABLE_P4 = False

class MetaMaterialParam():

    def __init__(self, name, desc, default, value):
        self.name = name
        self.value = value
        self.default = default
        self.desc = desc




def FindAreaResByMeshAndIndex(materialRes, mesh, areaIdx):
    geo = mesh.geometry
    meshName = mesh.name
    areaName = geo.GetMeshAreaName(mesh.meshIndex, areaIdx)
    if not materialRes:
        return 
    else:
        try:
            meshRes = materialRes.meshes[meshName]
        except KeyError:
            print ('KeyError! Missing mesh %s in materialRes' % meshName)
            return 
        else:
            try:
                areaRes = meshRes.areas[areaName]
            except KeyError:
                print ('KeyError! Missing area %s in meshRes %s' % (areaName, meshName))
                return 
        return areaRes



def ApplyMaterialStoreToEffect(effect, materialStore):
    params = effect.highLevelShader.parameterDescriptions
    for p in params:
        try:
            value = materialStore.FindParameter(p.parameterName)
            if value:
                effect.parameters[p.parameterName] = value
        except TypeError:
            print ('   Type Error copying %s in effect %s' % (p, effect.name))
            continue
        except KeyError:
            continue




class MaterialApplier():

    def __init__(self, mesh, materialRes):
        self.mesh = mesh
        self.materialRes = materialRes



    def Apply(self):
        geo = self.mesh.geometry
        meshIdx = self.mesh.meshIndex
        meshName = self.mesh.name
        try:
            meshRes = self.materialRes.meshes[meshName]
        except KeyError:
            print ('===== Did not find mesh %s in materialRes!' % meshName)
            return 
        else:
            for areaIdx in xrange(geo.GetMeshAreaCount(meshIdx)):
                areaName = geo.GetMeshAreaName(meshIdx, areaIdx)
                try:
                    areaRes = meshRes.areas[areaName]
                except KeyError:
                    print ('===== Did not find area %s in materialRes!' % areaName)
                    continue
                ApplyAreaRes(areaRes, self.mesh, areaIdx)

        self.mesh.BindLowLevelShaders([])




def LoadMaterialResFromObject(ob):
    resPath = GetResPath(ob)
    path = GetMaterialPathFromGeometryPath(resPath)
    materialRes = None
    if (path is not None):
        r = blue.ResFile()
        if r.FileExists(path):
            materialRes = trinity.Load(path, nonCached=True)
    return materialRes



def ApplyMaterialRes(ob, materialRes):
    if (materialRes == None):
        return 
    meshList = GetMeshList(ob)
    for mesh in meshList:
        try:
            meshRes = materialRes.meshes[mesh.name]
        except KeyError:
            print ('===== Did not find mesh %s in materialRes!' % mesh.name)
            continue
        if (mesh.geometry.isLoading or ((not mesh.geometry.isPrepared) or (not mesh.geometry.isGood))):
            mesh.AddGeometryPreparedCallback(MaterialApplier(mesh, materialRes).Apply)
        else:
            MaterialApplier(mesh, materialRes).Apply()




def LoadAndApplyMaterialRes(ob, materialRes = None):
    if (materialRes is None):
        materialRes = LoadMaterialResFromObject(ob)
    if (materialRes is None):
        materialRes = CreateMaterialResFromAreas(ob, '')
        return materialRes
    else:
        if (materialRes is None):
            return 
        ApplyMaterialRes(ob, materialRes)
        return materialRes



def SerializeMaterialStore(resPath, materialStore):
    import CCP_P4 as P4
    import os
    p4path = blue.rot.PathToFilename(resPath)
    destFolder = os.path.dirname(p4path)
    if not os.path.exists(destFolder):
        os.makedirs(destFolder)
    if (ENABLE_P4 is True):
        if not P4.P4FileInDepot(p4path):
            P4.P4Add(p4path)
        else:
            P4.P4Edit(p4path)
    blue.os.SaveObject(materialStore, resPath)



def SerializeMaterialRes(ob, materialRes):
    import CCP_P4 as P4
    import os
    resPath = GetResPath(ob)
    path = GetMaterialPathFromGeometryPath(resPath)
    if not path:
        return 
    p4path = blue.rot.PathToFilename(path)
    destFolder = os.path.dirname(p4path)
    if not os.path.exists(destFolder):
        os.makedirs(destFolder)
    if (ENABLE_P4 is True):
        if not P4.P4FileInDepot(p4path):
            P4.P4Add(p4path)
        else:
            P4.P4Edit(p4path)
    if not blue.os.SaveObject(materialRes, path):
        print ('===== SAVING FAILED for file %s =====' % p4path)



def CreateMaterialResFromAreas(ob, defaultMetatype):
    trinity.WaitForResourceLoads()
    meshes = GetMeshList(ob)
    if not meshes:
        return 
    else:
        materialRes = trinity.Tr2MaterialRes()
        for mesh in meshes:
            geo = mesh.geometry
            meshStore = trinity.Tr2MaterialMesh()
            materialRes.meshes[mesh.name] = meshStore
            for areaIdx in xrange(geo.GetMeshAreaCount(mesh.meshIndex)):
                areaName = geo.GetMeshAreaName(mesh.meshIndex, areaIdx)
                areaStore = trinity.Tr2MaterialArea()
                areaStore.metatype = str(defaultMetatype)
                areaStore.material = trinity.Tr2MaterialParameterStore()
                meshStore.areas[areaName] = areaStore
                params = GetParametersFromMeshArea(mesh, areaIdx, None)
                for p in params:
                    if (params[p].value is not None):
                        try:
                            if (params[p].value is not params[p].default):
                                areaStore.material.parameters[p] = params[p].value
                        except TypeError:
                            pass



        return materialRes



def SetAreaMetaShader(materialRes, mesh, areaIdx, metashader):
    geo = mesh.geometry
    meshName = mesh.name
    areaName = geo.GetMeshAreaName(mesh.meshIndex, areaIdx)
    try:
        meshRes = materialRes.meshes[meshName]
    except KeyError:
        print 'Mesh not found!'
        return 
    else:
        try:
            areaRes = meshRes.areas[areaName]
        except KeyError:
            print 'Area not found!'
            return 
    areaRes.metatype = str(metashader)
    ApplyAreaRes(areaRes, mesh, areaIdx)



def ApplyAreaRes(areaRes, mesh, areaIdx):
    metashaders = LoadMetaMaterialLibrary()
    try:
        currentShader = metashaders[str(areaRes.metatype)]
    except KeyError:
        print ('Unknown MetaMaterial: %s' % areaRes.metatype)
        return 
    else:
        for typeIdx in xrange(mesh.GetAreasCount()):
            areaList = mesh.GetAreas(typeIdx)
            if (areaList is not None):
                for area in areaList:
                    if (area.index == areaIdx):
                        areaList.remove(area)
                        break


        geo = mesh.geometry
        if not (geo.isPrepared and geo.isGood):
            print ('FAILED TO PREPARE MESH %s' % mesh.name)
            return 
        for area in currentShader:
            areaList = getattr(mesh, area['name'])
            if (areaList is not None):
                areaName = geo.GetMeshAreaName(mesh.meshIndex, areaIdx)
                newArea = trinity.Tr2MeshArea()
                newArea.index = areaIdx
                newArea.count = 1
                newArea.name = areaName
                newArea.effect = trinity.Tr2ShaderMaterial()
                newArea.effect.highLevelShaderName = area['shader']
                if ('situation' in area):
                    newArea.effect.defaultSituation = area['situation']
                newArea.effect.PopulateDefaultParameters()
                areaList.append(newArea)
                ApplyMaterialStoreToEffect(newArea.effect, areaRes.material)

    mesh.BindLowLevelShaders([])



def LoadMetaMaterialLibrary(resourceFile = METASHADER_PATH):
    global _gMetaShaderLibrary
    if (_gMetaShaderLibrary is not None):
        return _gMetaShaderLibrary
    else:
        _gMetaShaderLibrary = {}
        rf = blue.ResFile()
        if rf.FileExists(resourceFile):
            rf.Open(resourceFile)
            yamlStr = rf.read()
            rf.close()
            data = yaml.load(yamlStr)
            for metashader in data['metashaders']:
                _gMetaShaderLibrary[metashader['name']] = metashader['areas']

        return _gMetaShaderLibrary



def GetParametersFromMeshArea(mesh, areaIdx, areaRes):
    materialMap = {}
    for typeIdx in xrange(mesh.GetAreasCount()):
        areaList = mesh.GetAreas(typeIdx)
        if areaList:
            for area in areaList:
                if (area.index == areaIdx):
                    if (hasattr(area.effect, 'highLevelShader') and (area.effect.highLevelShader is not None)):
                        for desc in area.effect.highLevelShader.parameterDescriptions:
                            defVal = None
                            if (areaRes and (areaRes.material.parent and areaRes.material.parent.FindParameter(desc.parameterName))):
                                inheritedParam = areaRes.material.parent.FindParameter(desc.parameterName)
                            else:
                                inheritedParam = None
                            if inheritedParam:
                                defVal = inheritedParam
                            elif hasattr(desc, 'CreateDefaultParameter'):
                                defVal = desc.CreateDefaultParameter()
                            else:
                                defVal = None
                            materialMap[desc.parameterName] = MetaMaterialParam(name=desc.parameterName, desc=desc, value=defVal, default=defVal)

                        params = area.effect.parameters
                        for p in params:
                            materialMap[p].value = params[p]
                            if areaRes:
                                if ((p not in areaRes.material.parameters) and (p in materialMap)):
                                    materialMap[p].default = params[p]



    return materialMap



def GetMaterialPathFromGeometryPath(oldPath):
    import re
    newPath = re.sub('(?i)((.gr2)|(.red))$', '_Materials.red', oldPath)
    return newPath



def GetResPath(ob):
    if hasattr(ob, 'detailMeshes'):
        return ob.geometryResPath
    if (hasattr(ob, 'visualModel') and hasattr(ob.visualModel, 'geometryResPath')):
        return ob.visualModel.geometryResPath
    if hasattr(ob, 'placeableResPath'):
        return ob.placeableResPath



def GetMeshList(ob):
    try:
        if hasattr(ob, 'detailMeshes'):
            return ob.detailMeshes
        else:
            if hasattr(ob, 'visualModel'):
                return ob.visualModel.meshes
            if hasattr(ob, 'placeableRes'):
                return ob.placeableRes.visualModel.meshes
            return 
    except AttributeError:
        return 


exports = {'metaMaterials.ApplyAreaRes': ApplyAreaRes,
 'metaMaterials.ApplyMaterialRes': ApplyMaterialRes,
 'metaMaterials.LoadMaterialResFromObject': LoadMaterialResFromObject,
 'metaMaterials.LoadAndApplyMaterialRes': LoadAndApplyMaterialRes,
 'metaMaterials.LoadMetaMaterialLibrary': LoadMetaMaterialLibrary,
 'metaMaterials.FindAreaResByMeshAndIndex': FindAreaResByMeshAndIndex,
 'metaMaterials.SetAreaMetaShader': SetAreaMetaShader,
 'metaMaterials.SerializeMaterialRes': SerializeMaterialRes,
 'metaMaterials.SerializeMaterialStore': SerializeMaterialStore,
 'metaMaterials.GetMeshList': GetMeshList,
 'metaMaterials.GetParametersFromMeshArea': GetParametersFromMeshArea}

